'''
5.生成器-优化器（Generator-Optimizer） 定义：生成器（Generator）产生候选结果，优化器（Optimizer）对其进行评估、反馈、迭代改进。

应用场景：
代码生成与调试：生成代码 → 执行测试 → 修复错误 → 重新生成。
文案优化：生成广告文案 → 优化器评估“吸引力”“合规性” → 提出修改建议 → 重新生成。
对话系统：生成回复 → 评估“相关性”“安全性”“流畅度” → 调整输出。
AI绘画提示词优化：生成初始提示 → 根据图像质量反馈 → 优化提示词。
✅ 优势：实现闭环优化，提升输出质量，支持迭代改进。
''' 

# 导入必要的库和模块
from typing import Annotated, List
import operator
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage, SystemMessage
from typing_extensions import TypedDict, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from utils import init_llm, set_env, visualize_graph


class State(TypedDict):
    joke: str
    topic: str
    feedback: str
    funny_or_not: str


def llm_call_generator(state: State):
    """大模型生成笑话"""

    if state.get("feedback"):
        msg = llm.invoke(
            f"写一个关于 {state['topic']} 话题的笑话，但要考虑到反馈：{state['feedback']}"
        )
    else:
        msg = llm.invoke(f"写一个关于 {state['topic']} 话题的笑话")
    return {"joke": msg.content}


def llm_call_evaluator(state: State):
    """LLM 评估这个笑话"""

    grade = evaluator.invoke(f"对笑话进行评分 {state['joke']}")
    return {"funny_or_not": grade.grade, "feedback": grade.feedback}


def route_joke(state: State):
    """根据评估者的反馈返回笑话生成器或结束"""

    if state["funny_or_not"] == "funny":
        return "Accepted"
    elif state["funny_or_not"] == "not funny":
        return "Rejected + Feedback"




class Feedback(BaseModel):
    grade: Literal["funny", "not funny"] = Field(
        description="决定这个笑话是否有趣。",
    )
    feedback: str = Field(
        description="如果笑话不好笑，请提供有关如何改进的反馈。",
    )
    
def build_graph():
    """
    构建工作流图函数
    
    功能：
    1. 创建状态图实例
    2. 添加所有节点（编排器、工作节点、综合器）
    3. 定义节点之间的连接关系
    4. 编译并返回可执行的工作流图
    
    Returns:
        编译后的工作流图
    """
    
    # 创建状态图实例
    
    optimizer_builder = StateGraph(State)

    # 添加各个节点
    optimizer_builder.add_node("llm_call_generator", llm_call_generator)
    optimizer_builder.add_node("llm_call_evaluator", llm_call_evaluator)


    optimizer_builder.add_edge(START, "llm_call_generator")
    optimizer_builder.add_edge("llm_call_generator", "llm_call_evaluator")
    optimizer_builder.add_conditional_edges(
        "llm_call_evaluator",
        route_joke,
        {  
            "Accepted": END,
            "Rejected + Feedback": "llm_call_generator",
        },
    )

    optimizer_workflow = optimizer_builder.compile()

    return optimizer_workflow


if __name__ == "__main__":
    # 主程序入口
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    evaluator = llm.with_structured_output(Feedback)
    graph = build_graph()  # 构建工作流图
 
    # visualize_graph(graph)  # 可视化工作流图（已注释）

    # 执行工作流，生成关于LLM缩放定律的报告
    state = graph.invoke({"topic": "Cats"})
    print(state["joke"])  # 打印最终生成的报告







